#!/usr/bin/env python
# -*- coding=utf-8 -*-
import darknet
import os
import cv2
import numpy as np

CLASSES = {
    0: "chinese",
    1: "window",
    2: "grid",
    3: "headlight",
    4: "person",
    5: "plate",
    6: "top_window",
    7: "roof_rack",
    8: "mark",
    9: "tissue_box",
    10: "pendant",
    11: "spaer_wheel_rack",
    12: "visor",
}


if __name__ == '__main__':

    data_dir = '/home/xingwg/nas/data/Symbol/ehualu'
    list_file = '/home/xingwg/nas/data/Symbol/ehualu/image_list.txt'
    cfg_file = '/nfs-data/xingwg/deep_learning/darknet/models/mv_accessory_model/mvehicle_accessory.cfg'
    label_file = '/nfs-data/xingwg/deep_learning/darknet/models/mv_accessory_model/mvehicle_accessory.data'
    weight_file = '/nfs-data/xingwg/deep_learning/darknet/models/mv_accessory_model/mvehicle_accessory.weights'
    save_dir = "results"
    thresh = 0.25
    nms = 0.45
    hier_thresh = 0.5

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    if not os.path.exists(data_dir):
        print("data_dir not exist")
        exit(-1)

    if not os.path.isfile(list_file):
        print("list_file not exist")
        exit(-2)

    if not os.path.isfile(cfg_file):
        print("cfg_file not exist")
        exit(-3)

    if not os.path.isfile(weight_file):
        print("weight_file not exist")
        exit(-4)

    if not os.path.isfile(label_file):
        print("label_file not exist")
        exit(-5)

    label = darknet.load_meta(label_file.encode('utf-8'))

    net = darknet.load_net_custom(cfg_file.encode('utf-8'), weight_file.encode('utf-8'), 0, 1)

    input_w, input_h = (darknet.network_width(net), darknet.network_height(net))

    f = open(list_file, 'r')
    lines = f.readlines()
    f.close()

    f = open(os.path.join(save_dir, "darknet_yolov3.txt"), "w")
    for line in lines:
        image_name = line.strip()
        image_path = os.path.join(data_dir, image_name)
        if not os.path.isfile(image_path):
            print("image_path not exist -> {}".format(image_path))
            continue

        img = darknet.load_image(image_path.encode('utf-8'), 0, 0)

        img_w, img_h = img.w, img.h

        img = darknet.letterbox_image(img, input_w, input_h)

        darknet.image_write(img, os.path.join(save_dir, image_name).encode('utf-8'))

        darknet.network_predict(net, img.data)

        num = darknet.c_int(0)
        pnum = darknet.pointer(num)
        letter_box = 1
        detections = darknet.get_network_boxes(net, img_w, img_h, thresh, hier_thresh, None, 1, pnum, letter_box)

        num = pnum[0]
        darknet.do_nms_sort(detections, num, label.classes, nms)

        detect_results = list()
        for j in range(num):
            for i in range(label.classes):
                if detections[j].prob[i] > 0:
                    box = detections[j].bbox
                    tag = CLASSES[i]
                    x1 = int((box.x - box.w * 0.5) * img_w)
                    y1 = int((box.y - box.h * 0.5) * img_h)
                    x2 = int((box.x + box.w * 0.5) * img_w)
                    y2 = int((box.y + box.h * 0.5) * img_h)

                    x1 = 0 if x1 < 0 else x1
                    y1 = 0 if y1 < 0 else y1
                    x2 = img_w - 1 if x2 >= img_w else x2
                    y2 = img_h - 1 if y2 >= img_h else y2

                    detect_results.append((tag, detections[j].prob[i], (x1, y1, x2, y2)))

        detect_results = sorted(detect_results, key=lambda x: -x[1], reverse=True)

        darknet.free_detections(detections, num)

        for det in detect_results:
            text = "{} {} {} {} {} {} {}".format(image_name, det[0], det[1], det[2][0], det[2][1], det[2][2], det[2][3])
            print(text)
            f.write(text)
            f.write("\n")

    f.close()
